tensorflow gpu

tensorflow默认抢占服务器所有GPU的所有显存，一段小程序也会占用所有GPU资源。

解决方案：

为tf指定GPU：

import os
os.environ[‘CUDA_VISIBLE_DEVICE’] = ‘1,2'
# os.environ['CUDA_VISIBLE_DEVICES'] = ''  # 设置为空，则使用cpu。或者='-1'

或者

import tensorflow as tf
with tf.device('/gpu:0'):       # Run nodes with GPU 0
    m1 = tf.constant([[3, 5]])
    m2 = tf.constant([[2],[4]])
    product = tf.matmul(m1, m2)    

sess = tf.Session()
print(sess.run(product))

sess.close()

或者设置环境变量

$ export CUDA_VISIBLE_DEVICES=0,1
# 指定cpu，指定 visible devices为空
$ export CUDA_VISIBLE_DEVICES=

其次是限制在每个可用GPU上占据需要使用的内存大小：

gpu_options = tf.GPUOptions(allow_growth=True)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

该设置会启用最少的GPU显存来运行程序。

在tensorflow中定义session时作如下设置，该设置会强制程序只占用指定比例的GPU显存。

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.4 # 占用GPU40%的显存
session = tf.Session(config=config)

多GPU

import tensorflow as tf

c = []
for i, d in enumerate(['/gpu:0', '/gpu:1', '/gpu:2']):
    with tf.device(d):
        a = tf.get_variable(f"a_{i}", [2, 3], initializer=tf.random_uniform_initializer(-1, 1))
        b = tf.get_variable(f"b_{i}", [3, 2], initializer=tf.random_uniform_initializer(-1, 1))
        c.append(tf.matmul(a, b))

with tf.device('/cpu:0'):
    sum = tf.add_n(c)

sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))

init = tf.global_variables_initializer()
sess.run(init)

print(sess.run(sum))
# [[-0.36499196 -0.07454088]
# [-0.33966339  0.30250686]]

如何设置多个worker gpu

tensor2tensor中的worker-gpu配置，是怎样实现的。